import os
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
# from scipy.misc import imread # which has delete the imread
import imageio




def computeMean():
    R_channel = 0
    G_channel = 0
    B_channel = 0
    train = []
    test = []
    with open('D:\\data\\jindi\\SAVE\\JinDiTrain.txt', 'r') as f:
        while True:
            content = f.readline()
            if not content:
                break
            train.append(content)

    with open('D:\\data\\jindi\\SAVE\\JinDiTest.txt', 'r') as f:
        while True:
            content = f.readline()
            if not content:
                break
            test.append(content)
    result = train + test

    for idx in result:
        # idx = idx.replace('\n','') #.split(' ')
        idx = idx[0:-3]
        print(idx)
        img = imageio.imread(idx)  / 65536.0
        img = np.array(img)
        if img.ndim != 3:
            img = np.expand_dims(img, axis= 2)
            img = np.repeat(img,3,axis=2)
        (w,h,c) = img.shape

        R_channel = R_channel + np.sum(img[:, :, 0]) / (w*h)
        G_channel = G_channel + np.sum(img[:, :, 1]) / (w*h)
        B_channel = B_channel + np.sum(img[:, :, 2]) / (w*h)

    num = len(result)  # 这里（512,512）是每幅图片的大小，所有图片尺寸都一样
    R_mean = R_channel / num
    G_mean = G_channel / num
    B_mean = B_channel / num
    return R_mean, G_mean, B_mean


def computeStd( R_mean, G_mean, B_mean):
    R_channel = 0
    G_channel = 0
    B_channel = 0

    train = []
    test = []
    with open('D:\\data\\jindi\\SAVE\\JinDiTrain.txt', 'r') as f:
        while True:
            content = f.readline()
            if not content:
                break
            train.append(content)

    with open('D:\\data\\jindi\\SAVE\\JinDiTest.txt', 'r') as f:
        while True:
            content = f.readline()
            if not content:
                break
            test.append(content)
    result = train + test

    for idx in result:
        # idx = idx.replace('\n', '').split(' ')
        idx = idx[0:-3]
        img = imageio.imread(idx) / 65536.0

        if img.ndim != 3:
            img = np.expand_dims(img, axis=2)
            img = np.repeat(img, 3, axis=2)
        (w, h, c) = img.shape
        R_channel = R_channel + np.sum((img[:, :, 0] - R_mean) ** 2) / (w*h)
        G_channel = G_channel + np.sum((img[:, :, 1] - G_mean) ** 2)  / (w*h)
        B_channel = B_channel + np.sum((img[:, :, 2] - B_mean) ** 2) / (w*h)
    num = len(result)
    R_var = np.sqrt(R_channel / num)
    G_var = np.sqrt(G_channel / num)
    B_var = np.sqrt(B_channel / num)
    return R_var, G_var, B_var


if __name__ == '__main__':
    print('compute mean and std')
    filepath = 'D:/data/data/train/mine'  # 数据集目录
    filepath = 'D:\\data\\jindi\\SAVE\\mine'
    mine = os.listdir(filepath)

    filepath = 'D:\\data\\jindi\\SAVE\\waste'
   #  filepath = ''
    waste = os.listdir(filepath)
    pathDir = mine + waste
    # pathDir
    R_mean, G_mean, B_mean = computeMean()
    print('compute std result! ')
    R_var, G_var, B_var = computeStd(R_mean, G_mean, B_mean)
    print("R_mean is %f, G_mean is %f, B_mean is %f" % (R_mean, G_mean, B_mean))
    print("R_var is %f, G_var is %f, B_var is %f" % (R_var, G_var, B_var))
